#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
INT8 量化模型评测脚本
评估量化模型的性能：准确率、推理速度、显存占用
"""
import os
import json
import ast
import sys
import time
from typing import List, Dict, Optional

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 使用量化推理模块
from utils.inference_quantized import generate_code


def _first_existing(paths: List[str]) -> Optional[str]:
    """返回路径列表中第一个存在的路径"""
    for p in paths:
        if os.path.exists(p):
            return p
    return None


def _detect_codesearchnet_paths():
    """检测处理后的 CodeSearchNet-Python 测试集路径
    
    Returns:
        (test_path, test_candidates): 测试集路径和候选路径列表
    """
    test_candidates = [
        os.path.join('data', 'processed', 'codesearchnet_test_text2code.jsonl'),
        os.path.join('data', 'processed', 'complex_test_text2code.jsonl'),
    ]
    return _first_existing(test_candidates), test_candidates


def _read_jsonl(path: str) -> List[Dict]:
    """读取 JSONL 格式文件"""
    data = []
    if not os.path.exists(path):
        return data
    with open(path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                data.append(json.loads(line))
            except json.JSONDecodeError:
                continue
    return data


def _is_valid_python_function(code: str) -> bool:
    """检查代码是否为语法正确的 Python 函数"""
    s = (code or '').strip()
    if not s.startswith('def '):
        return False
    try:
        tree = ast.parse(s)
    except SyntaxError:
        return False
    fdefs = [n for n in tree.body if isinstance(n, ast.FunctionDef)]
    return len(fdefs) >= 1


def _rough_quality_check(prompt: str, code: str) -> bool:
    """基于关键词和 AST 结构进行粗略质量检查
    
    Args:
        prompt: 输入提示
        code: 生成的代码
    
    Returns:
        bool: 是否通过检查
    """
    s = (code or '').strip().lower()
    p = (prompt or '').strip().lower()
    
    if not _is_valid_python_function(code):
        return False
    
    # 排除明显错误的模式
    if 'class ' in s:
        if 'function' in p and 'class' not in p:
            return False
    
    # 关键词匹配检查
    if 'return' not in s:
        if 'print' not in p:
            return False
    
    return True


def main(n: int | None = None, model: str = "auto"):
    """主评测函数
    
    Args:
        n: 评测样本数量，None 表示全部
        model: 模型类型 (auto/complex/multitask)
    """
    # 设置量化模型环境变量
    os.environ["T2C_QUANTIZE_MODEL"] = model
    
    if model == "multitask":
        print("[INFO] 量化多任务 LoRA 模型")
        model_name = "Multitask"
        original_acc = 0.75  # 多任务基线
    elif model == "complex":
        print("[INFO] 量化复杂任务 LoRA 模型")
        model_name = "Complex"
        original_acc = 0.70  # 复杂任务基线
    else:  # auto
        print("[INFO] 自动选择最佳可用模型")
        model_name = "Auto-selected"
        original_acc = 0.70  # 保守基线
    
    test_path, test_candidates = _detect_codesearchnet_paths()
    test_records = _read_jsonl(test_path) if test_path else []
    
    if not test_records:
        looked = ", ".join(test_candidates)
        print(f"[ERROR] 未找到处理后的 CodeSearchNet-Python 测试集。查找路径: {looked}\n"
              f"请先运行预处理脚本 (如 preprocess_data.py) 生成数据集\n"
              f"提示: 当前复杂任务数据集保存为 'data/processed/complex_*_text2code.jsonl'")
        return 1

    total = 0
    pass_cnt = 0
    total_time = 0.0
    
    print("\n" + "=" * 70)
    print(f"   INT8 量化模型评估 - {model_name}")
    print("=" * 70)
    print(f"评测样本数: {n if n else len(test_records)}\n")
    
    for item in (test_records[:n] if n else test_records):
        instr = item.get('instruction', '')
        total += 1
        
        # 测量推理时间
        start_time = time.time()
        code = generate_code(instr, mode="complex")
        elapsed = time.time() - start_time
        total_time += elapsed
        
        # 使用粗略规则检查
        ok = _rough_quality_check(instr, code)
        
        if ok:
            pass_cnt += 1
        
        status = 'PASS' if ok else 'FAIL'
        print(f"[{total:3d}] {status} ({elapsed:.2f}s)")

    acc = pass_cnt / total if total else 0.0
    avg_time = total_time / total if total else 0.0
    
    print("\n" + "=" * 70)
    print(f"   评测结果 - {model_name} INT8 量化")
    print("=" * 70)
    print(f"Total samples:        {total}")
    print(f"Passed:               {pass_cnt}")
    print(f"Failed:               {total - pass_cnt}")
    print(f"Accuracy:             {acc:.1%} ({pass_cnt}/{total})")
    print(f"Average inference:    {avg_time:.2f}s per sample")
    print(f"Total time:           {total_time:.1f}s")
    print("=" * 70)
    
    # 与原始模型对比
    performance_retention = (acc / original_acc) * 100 if original_acc > 0 else 0
    
    print("\n" + "=" * 70)
    print("   性能对比")
    print("=" * 70)
    print(f"Original {model_name}:    {original_acc:.1%} accuracy")
    print(f"INT8 quantized:       {acc:.1%} accuracy")
    print(f"Performance retention: {performance_retention:.1f}%")
    
    if acc >= 0.67:
        print("\n[SUCCESS] 量化性能优秀！")
        print("   性能保持 > 95%，可部署使用。")
    elif acc >= 0.60:
        print("\n[WARNING] 检测到性能下降。")
        print("   建议使用动态量化或原始模型。")
    else:
        print("\n[FAIL] 性能损失显著。")
        print("   INT8 量化不适合此任务。")
    
    print("=" * 70 + "\n")
    
    return 0


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description="INT8 量化模型评测脚本")
    parser.add_argument('--n', type=int, default=None, help="评测样本数量 (默认: 全部)")
    parser.add_argument('--model', type=str, default="auto", 
                        choices=["auto", "complex", "multitask"],
                        help="量化和评测的模型类型 (默认: auto)")
    args = parser.parse_args()
    raise SystemExit(main(args.n, args.model))
